#include "caffe_insight.h"

CaffeInsight::CaffeInsight(const std::string& modelFile, const std::string& weightFile, const std::string& mean_file, const std::string& mean_value, int batchsize, int platform)
{
  this->modelFile = modelFile;
  this->weightFile = weightFile;
  this->batchsize = batchsize;
  this->platform = platform;

  switch (platform)
  {
    case 0:
      caffe::Caffe::set_mode(caffe::Caffe::GPU);
      break;
    case 1:
      caffe::Caffe::set_mode(caffe::Caffe::CPU);
      break;
    default:
      caffe::Caffe::set_mode(caffe::Caffe::GPU);
      break;
  }

  this->net.reset(new caffe::Net<float>(modelFile, caffe::TEST));
  this->net->CopyTrainedLayersFrom(weightFile);
  // CHECK_EQ(m_net->num_inputs(), 1) << "Network should have exactly one input.";
  // CHECK_EQ(m_net->num_outputs(), 1) << "Network should have exactly one output.";
  this->inputBlob = this->net->input_blobs()[0];
  this->ichannels = this->inputBlob->channels();
  this->isize = cv::Size(this->inputBlob->width(), this->inputBlob->height());

  this->LoadMeanMat(mean_file, mean_value);

  this->inputBlob->Reshape(batchsize, this->ichannels, this->isize.height, this->isize.width);
  this->net->Reshape(); //NOTE: Forward dimension change to all layers.
}

CaffeInsight::~CaffeInsight()
{
}

void CaffeInsight::LoadMeanMat(const std::string& mean_file, const std::string& mean_value)
{
  cv::Scalar channel_mean;
  if (!mean_file.empty()) {
    CHECK(mean_value.empty()) << "[Error]: Can't specify mean_file and mean_value at the same time!";
    caffe::BlobProto blob_proto;
    caffe::ReadProtoFromBinaryFileOrDie(mean_file.c_str(), &blob_proto);

    //Convert from BlobProto to Blob<float>
    caffe::Blob<float> mean_blob;
    mean_blob.FromProto(blob_proto);
    CHECK_EQ(mean_blob.channels(), this->ichannels) << "[Error]: Number of channels of mean file doesn't match input layer.";

    std::vector<cv::Mat> channels;
    float* data = mean_blob.mutable_cpu_data();
    for (int i = 0; i < this->ichannels; i++) {
      //Extract an individual channel.
      cv::Mat channel(mean_blob.height(), mean_blob.width(), CV_32FC1, data);
      channels.push_back(channel);
      data += mean_blob.height() * mean_blob.width();
    }

    //Merge the separate channels into a single image.
    cv::Mat mean;
    cv::merge(channels, mean);

    //Compute the global mean pixel value and create a mean image filled with this value.
    channel_mean = cv::mean(mean);
    this->meanMat = cv::Mat(this->isize, mean.type(), channel_mean);
  }
  if (!mean_value.empty()) {
    CHECK(mean_file.empty()) << "[Error]: Can't specify mean_file and mean_value at the same time!";
    std::stringstream ss(mean_value);
    std::vector<float> values;
    std::string item;
    while (getline(ss, item, ',')) {
      float value = std::atof(item.c_str());
      values.push_back(value);
    }
    CHECK((values.size() == 1) || (values.size() == this->ichannels)) << "[Error]: Specify either 1 mean_value or as many as channels: " << this->ichannels;

    std::vector<cv::Mat> channels;
    for (int i = 0; i < this->ichannels; i++) {
      //Extract an individual channel.
      cv::Mat channel(this->isize.height, this->isize.width, CV_32FC1, cv::Scalar(values[i]));
      channels.push_back(channel);
    }
    cv::merge(channels,this->meanMat);
  }
}

void CaffeInsight::WrapInputBlob(std::vector<std::vector<cv::Mat>>& images, int batchsize)
{
  CHECK(batchsize == this->batchsize) << "[Error]: Your inputblob's batchsize is not same as the detector's!";

  this->inputBlob = this->net->input_blobs()[0];
  int width = this->inputBlob->width();
  int height = this->inputBlob->height();
  float* dataptr = this->inputBlob->mutable_cpu_data();

  for (int i = 0; i < batchsize; i++)
  {
    std::vector<cv::Mat>* input_channels = &images[i];
    for (int j = 0; j < this->inputBlob->channels(); j++)
    {
      cv::Mat channel(height, width, CV_32FC1, dataptr);
      input_channels->push_back(channel);
      dataptr += width * height;
    }
  }
}

void CaffeInsight::Preprocess(const cv::Mat* images, int batchsize, std::vector<std::vector<cv::Mat>>& inputBlob)
{
  if (batchsize != inputBlob.size()) {
    std::cerr << "[Error]: input batchsize isn't match!" << std::endl;
    return;
  }
  for (int i = 0; i < batchsize; i++) {
    std::vector<cv::Mat>* input_channels = &(inputBlob[i]);
    Preprocess(images[i], input_channels);
  }
}

void CaffeInsight::Preprocess(const cv::Mat& img, std::vector<cv::Mat>* input_channels)
{
  //Convert the input image to the input image format of the network.
  cv::Mat sample;
  if (img.channels() == 3 && this->ichannels == 1) {
    cv::cvtColor(img, sample, cv::COLOR_BGR2GRAY);
  }
  else if (img.channels() == 4 && this->ichannels == 1) {
    cv::cvtColor(img, sample, cv::COLOR_BGRA2GRAY);
  }
  else if (img.channels() == 4 && this->ichannels == 3) {
    cv::cvtColor(img, sample, cv::COLOR_BGRA2BGR);
  }
  else if (img.channels() == 1 && this->ichannels == 3) {
    cv::cvtColor(img, sample, cv::COLOR_GRAY2BGR);
  }
  else {
    sample = img;
  }

  cv::Mat sample_resized;
  if (sample.size() != this->isize) {
    cv::resize(sample, sample_resized, this->isize);
  }
  else {
    sample_resized = sample;
  }

  cv::Mat sample_float;
  if (this->ichannels == 3) {
    sample_resized.convertTo(sample_float, CV_32FC3, 0.0078125, -1);
  }
  else {
    sample_resized.convertTo(sample_float, CV_32FC1, 0.0078125, -1);
  }

  cv::Mat sample_normalized;
  cv::subtract(sample_float, this->meanMat, sample_normalized);

  //This operation will write the separate BGR planes directly to the input layer of the network
  //because it is wrapped by the cv::Mat objects in input_channels.
  cv::split(sample_normalized, *input_channels);
  //CHECK(reinterpret_cast<float*>(input_channels->at(0).data) == this->net->input_blobs()[0]->cpu_data()) << "[Error]:Input channels are not wrapping the input layer of the network.";
}

void CaffeInsight::Predict(const cv::Mat* images, float* result, int batchsize, const char* output_layer)
{
  std::vector<std::vector<cv::Mat>> inputBlob;
  inputBlob.resize(batchsize);
  WrapInputBlob(inputBlob, batchsize);

  Preprocess(images, batchsize, inputBlob);

  this->net->Forward(); //this->net->ForwardPrefilled();
  const boost::shared_ptr<caffe::Blob<float> > output = this->net->blob_by_name(output_layer);

  if (result != NULL) {
    memcpy(result, output->cpu_data(), output->count()*sizeof(float));
  }
  return;
}

